"""Module to implement all flux functions.

A flux function is responsible for determining the flux at a cell
interface. This face flux can then be used in an appropriate
conservative update formula involving face flux values.

"""

from hogs.solvers.limiters import waf_limiters
from hogs.grids.reconstruct1d import ReconstructionType, MUSCL
from hogs.solvers.limiters import limiters

from optparse import OptionGroup
import numpy 

class FluxFunction(object):
    """Base class for all flux functions."""

    def set_grid(self, grid):
        """Set the grid object for the fluc function"""
        self.grid = grid

    def compute_flux(self):
        raise NotImplementedError("FluxFunction::compute_flux called")

    def get_options(self, parser):
        """Get flux function specific options"""
        pass

    def set_options(self, options):
        """Set Flux specific options"""
        pass

    def _reconstruct(self):
        """Reconstruct the variables in the cells """


###########################################################################
#`AdvectionFlux1D`
###########################################################################
class AdvectionFlux1D(FluxFunction):
    r"""Flux for the linear advection equation.

    .. attribute:: a

      Constant speed for advection.
    

    For the linear advection equation in the form:
    :math:`\frac{\partial Q}{\partial t} + a \frac{\partial
    Q}{\partial x} = 0`, the flux at the face :math:`i + \frac{1}{2}`
    is given by:

    .. math::

        F_{i + \frac{1}{2}} = a Q_{i} + a^{-}( Q_{i+1} - Q_{i} ) -
        \frac{1}{2}(a^+ - a^-)\left( 1 - \frac{\Delta t}{\Delta x}
        (a^+ - a^-) \right ) \tilde{W}_{i+\frac{1}{2}}

    where the limited wave :math:`\tilde{W}` is given by:

    .. math::

        \tilde{W}_{i+\frac{1}{2}} = (Q_{i + 1} -
        Q_{i})\,\phi(\theta_{i + \frac{1}{2}})

    Here, :math:`\phi` is an appropriate limiter function and
    :math:`\theta` is given as:

    .. math::

        \theta_{i + \frac{1}{2}} = \frac{Q_{I + \frac{1}{2}}}{Q_{i +
        \frac{1}{2}}}


        with\,\,  I = \begin{cases}
              i - 1 & if \lambda^p > 0 \\
              i + 1 & if \lambda^p < 0
              \end{cases}


    """
    def __init__(self, a):
        self.a = a

    def compute_flux(self, grid, dt):
        """Compute the flux for each face of the grid."""
        faces = grid.faces
        a = self.a

        dx = grid.dx; k = dt/dx

        # get the number of boundary cells
        nb = grid.nb

        for face in grid.faces[nb:-nb]:

            ql, qr = face.get_conserved_variables()

            f = face.flux

            # compute the first order term
            f[0] = a * ql[0] + min(a, 0) * ( qr[0] - ql[0] )

###########################################################################
#`LinearSolverFlux`
###########################################################################
class LinearSolverFlux(FluxFunction):
    r"""Flux function for the linear solver.

    The flux for a face at :math:`i + \frac{1}{2}` is given by:

    .. math::

        F_{i + \frac{1}{2}} = AQ_{i} + \sum_{p=1}^{m} (\lambda^p)^-
        W_{i+\frac{1}{2}}^p + \frac{1}{2}\sum_{p=1}^{m}\left(
        |\lambda^p| (1 - \frac{\Delta t}{\Delta
        x}|\lambda^p|)\right)\tilde{W}_{i + \frac{1}{2}^p}

    To implement the flux in this way, wach subclass must define
    functions to return the wave strengths :math:`\alpha`, the
    eigenvalues :math:`\lambda` and the corresponding eigenvectors
    :math:`r`.

    In addition, each sub class should implement two functions:

    *compute_first_order_flux*:
        This function should compute an appropriate first order flux
        (eg. Godunov's flux) and store the contribution in each face's
        flux attribute.

    *compute_high_order_flux*:
        This function should compute the high order correction terms
        for the flux (using a slope limiter approach perhaps) and append
        the flux contribution to the face's flux attribute.

    The default limiter function for the LinearSolverFlux is the
    UpwindLimiter which results in a first order method (that is, no
    correction is applied).

    Available limiters are defined in :file:`limiters.py`

    """
    def __init__(self, limiter_type=0):
        self.defaults = {"limiter_type":limiter_type}
    
    def get_options(self, parser):
        """General linear solver options"""

        opt = OptionGroup(parser, "Limiter options for the linear solver")

        opt.add_option("--limiter", action="store", type="int",
                       dest="limiter_type",
                       default = self.defaults["limiter_type"],
                       help="The limiter function to use:\t\t\t\t" + \
                       "".join(['%d - %-51s'%(d,s[0]) for d,s in
                                enumerate(limiters.limiter_functions)]))
        return opt

    def set_options(self, options):
        """Set any specific options"""
        limiter_type = options.get("limiter_type")
        if limiter_type is not None:
            self.limiter_func = limiters.limiter_functions[limiter_type][1]()
    
    def compute_flux(self, grid, dt):
        """Compute the flux for each face."""

        nb = grid.nb

        for face in grid.faces:
            if face.id == FaceType.Interior:

                # compute the first order flux
                self.compute_first_order_flux(face, grid, dt)

                # compute high order correction terms
                self.compute_high_order_flux(face, grid, dt)            

#########################################################################
# `WAFFlux` class
#########################################################################
class WAFFlux(FluxFunction):
    r"""Weighted average flux of Toro (Chapter 13)

    The WAF method defines the flux at the interface :math:`i +
    \frac{1}{2}` as :

    .. math::

      F_{i + \frac{1}{2}} = \frac{1}{2}( F_i + F_{i+1} ) -
      \frac{1}{2}\sum_{p = 1}^{p=m} sign(c_p) \phi_{i + \frac{1}{2}}^m
      \Delta F_{i + \frac{1}{2}}^p

    Here :math:`m` is the number of waves in the solution to an
    Riemann problem. :math:`\phi` is a limiter function for each wave.

    """            
    def __init__(self, limiter_type=0):
        try:
            self.defaults.update( dict( limiter_type=limiter_type ) )
        except AttributeError:
            self.defaults = dict( limiter_type=limiter_type )
        
    def get_options(self, parser):
        r"""Get options for the WAF flux. 

        The main options for this flux is the limiter function for the
        WAF flux which are defined in
        :file:`hogs/solvers/limiters/waf_limiters.py` and can be
        one of

           * SuperBee
           * vanLeer
           * vanAlbada
           * MinBee

        """
        opt = OptionGroup(parser, "Limiter functions for the WAF flux")

        opt.add_option("--limiter", action="store", type="int",
                       dest="limiter_type",
                       default = self.defaults["limiter_type"],
                       help="The limiter function to use:\t\t\t\t" + \
                       "".join(['%d - %-51s'%(d,s[0]) for d,s in
                                enumerate(waf_limiters.limiter_functions)]))
        return opt

    def set_options(self, options):
        """Set options for the WAF flux"""

        limiter_type = options.get("limiter_type")
        self.limiter_func = waf_limiters.limiter_functions[limiter_type][1]()

    def compute_flux(self, grid, dt):
        """Compute the WAF flux"""
        nb = grid.nb
        k = dt/grid.dx

        # For the WAF flux, we need to sample the data states between
        # the waves in the solution to the Riemann problem. For the
        # Euler equations in 1D, we have 3 waves and 4 states. For the
        # Burger's equation in 1D, we have 1 wave and hence two states.
        U = self._get_data_state_vector()

        # In addition to the data states, we need the wave speed
        # estimates for the waves in the Riemann solution.
        S = self._get_wave_speed_vector()

        # compute the flux for each face
        for face in grid.faces[nb:-nb]:

            # get the states on either side of the discontinuity
            ql, qr = face.get_conserved_variables()
            
            # for the flux computation, we need a quantity that is
            # known to jump across each wave. This will help us in
            # defining the limiter function for the WAF flux.
            dq = self._get_jump(ql, qr)

            # compute the data states 'U' and wave speeds 'S'
            self._compute_waf_states(ql, qr, face, U, S)

            # compute the first order unstable flux
            f = self._compute_unstable_flux(ql, qr)

            # get the neighboring jumps for the limiter function
            ql, qr = face.get_conserved_variables(i=1)
            dqp1 = self._get_jump(ql, qr)

            ql, qr = face.get_conserved_variables(i=-1)
            dqm1 = self._get_jump(ql, qr)

            # second order corrections or characteristic wise limiting
            for i in range(self.rsolver.nwaves):

                # Courant number for this wave
                ck = S[i] * k

                # numerical check if the jump is small. rk is chosen
                # in this case to be -1 since all limiter functions
                # will default to the Godunov's method in this case.
                if abs(dq) < 1e-10:
                    rk = -1.0

                else:
                    if ck > 0:
                        rk = dqp1/dq
                    else:
                        rk = dqm1/dq

                # limiting constant based on the flow parameter rk
                phik = self.limiter_func.limit(rk, abs(ck))                
            
                # jump in the k-th characteristic: F[U(i+1)] - F[U(i)]
                delF = self._compute_delF(U,i)

                # add the correction term to the flux (equation 14.23)
                f -= 0.5 * numpy.sign(ck) * phik * delF

            # store the flux for this face
            face.flux = f

    def _get_data_state_vector(self):
        raise NotImplementedError("WAFFlux::_get_data_state_vector called")

    def _get_wave_speed_vector(self):
        raise NotImplementedError("WAFFlux::_get_wave_speed_vector called")
    
    def _get_jump(self, ql, qr):
        raise NotImplementedError("WAFFlux::_get_jump called")

    def _compute_unstable_flux(self, ql, qr):
        raise NotImplementedError("WAFFlux::_compute_unstable_flux called")

    def _compute_delF(self, U, i):
        raise NotImplementedError("WAFFlux::_compute_delF called")

    def _compute_waf_states(self, ql, qr, face, U, S):
        raise NotImplementedError("WAFFlux::_compute_waf_states called")

    def _compute_delF(self, U, i):
        raise NotImplementedError("WAFFlux::_compute_delF called")

###########################################################################
# `LxfFlux` class
###########################################################################
class LxfFlux(FluxFunction):
    """General Lax Friedrichs flux function

    The Lax Friedrichs flux (lxf) at the interface :math:`i +
    \frac{1}{2}` is given by:

    .. math::

      F_{i + \frac{1}{2}} = \frac{1}{2}(F_{i} + F_{i + 1})
      -\frac{1}{2} a(Q_{i+1} - Q_{i})

    """
    def __init__(self, muscl=0, ppm=-1, eno=-1):
        self.defaults = dict( muscl=muscl, ppm=ppm, eno=eno )

    def get_options(self, parser):
        """Get options for the AcousticswithAdvectionLxfFlux1D"""

        opt = OptionGroup(parser, "Options for the reconstruction method")

        opt.add_option("--muscl", action="store", type="int", dest="muscl",
                       default=self.defaults.get("muscl"),
                       help="MUSCL reconstruction scheme")
                       
        opt.add_option("--eno", action="store", type="int", dest="eno",
                       default=self.defaults.get("eno"),
                       help="ENO reconstruction")

        opt.add_option("--ppm", action="store", type="int", dest="ppm",
                       default=self.defaults.get("ppm"),
                       help="PPM reconstruction")
        
        return opt

    def set_options(self, options):
        """Set options for the AcousticswithAdvectionLxfFlux1D"""
        options = options or self.defaults

        # set the reconstruction options
        muscl = options.get("muscl")
        ppm = options.get("ppm")
        eno = options.get("eno")

        # check the reconstruction method
        self._check_reconstruction_options(muscl, ppm, eno)

        # set the MUSCL related options
        if self.reconstruction_method == ReconstructionType.MUSCL:

            # Unhash to actually implement this
            #limiter = options.get("muscl_limiter")
            self.grid.set_reconstruction_function(self.reconstruction_method)

    def compute_flux(self, grid, dt):
        """Compute flux for each face."""
        grid = self.grid
        k = grid.dx/dt
        
        for face in grid.faces[grid.nb:-grid.nb]:

            # get the conserved variables on either side of the face
            ql, qr = face.get_conserved_variables(i=0)

            # compute the unstable flux
            f = self._compute_unstable_flux(ql, qr)

            # get the max sound speed for the interface
            c = self._compute_max_sound_speed()
            
            # compute the diffusive term
            f -= 0.5 * c * k *( qr - ql )

            # store the flux
            face.flux[:] = f        
        
    def _check_reconstruction_options(self, muscl, ppm, eno):
        if muscl >= 0:
            msg = "MUSCL type reconstruction selected"
            assert ppm < 0, msg
            assert eno < 0, msg

            self.reconstruction_method = ReconstructionType.MUSCL

        elif ppm >= 0:
            msg = "PPM type reconstruction selected"
            assert muscl < 0, msg
            assert eno < 0, msg

            self.reconstruction_method = ReconstructionType.PPM

        elif eno > 1:
            msg = "ENO type reconstruction selected"
            assert muscl < 0, msg
            assert ppm < 0, msg

            self.reconstruction_method = ReconstructionType.ENO
